import json
import time

import gradio as gr
import numpy as np
from PIL import Image

from taskinfo.Uie import Uie
from taskinfo.uv_pipeline import UvPipeline

uv_predict = UvPipeline(model_path=r"D:\ocr\onnx\deploy\uvdoc\inference.onnx", use_openvino=True)

uie = Uie(
    uie_model_path=r"D:\ocr\onnx\deploy\rong_cheng\information_extraction\static\inference.onnx",
    det_model_path=r"D:\ocr\onnx\deploy\text_system\det_1.0.0\inference.onnx",
    rec_model_path=r"D:\ocr\onnx\deploy\text_system\rec_1.0.0\inference.onnx",
    use_openvino=True
)


def predict(image: Image.Image, text: str):
    image_data = np.array(image)
    result = uv_predict(image_data)
    if result['output_imgs']:
        image_data = result['output_imgs'][0]

    img = Image.fromarray(image_data)
    r_channel = img.split()[0]
    rgb_image = r_channel.convert('RGB')
    image_data = np.array(rgb_image)
    # input_data = {"original_image": image_data, "image": image_data.copy()}
    processed_image = rgb_image
    uie.set_schema_list(text.replace("，", ",").split(","))
    start_time = time.process_time()
    text_tesult = uie(image_data.copy())
    print(f"耗时：{time.process_time() - start_time}")
    return processed_image, json.dumps(text_tesult, ensure_ascii=False, indent=4)


def gradio_interface(image, text):
    return predict(image, text)


iface = gr.Interface(
    fn=gradio_interface,
    inputs=["image", gr.Textbox(placeholder="提取文本以,分割")],
    outputs=['image', 'text'],
    title="信息提取demo",
    description="上传一张图片，提取关键信息"
)

if __name__ == '__main__':
    iface.launch(server_name="0.0.0.0")
